So in the end, this is a softmax. @rkhosrowshahi I think this means that moco is essentially just a composition with this weighting, where essentially you give yy_t to it, and then multiply yy_t by the obtained weights. Is that correct? If yes, I think we should change #676 accordingly.

Can remove the description of what MoDo is, or fix it (it's not doable with autogram with the changes I suggest).

Can't inherit from _GramianWeighting anymore with the changes I propose. Need to either inherit from _MatrixWeighting, or inherit from Weighting[Matrix] and override call to make a more specific docstring (e.g. explaining exactly what the matrix is: J_1 @ J_2^T)

Also need to fix the main docstring accordingly.

Arguably the whole method (given in the usage example) is called Modo. Can remove this last part. The part explaining the acronym (in parentheses) is nice though.

Not a gramian anymore with the changes i propose.

This is a bit wrong, we can't talk about a Gramian there and it should itself come from 2 batches. The last part about autogram should be removed.

I would simplify this explanation to make it clear that this is just doing basic MoDo, so that people actually follow this usage example if they wanna reproduce MoDo.

Something like:
Train a model with MoDo.

The role of their paper is to explain what MoDo is. The role of TorchJD is to make it extremely easy to reproduce MoDo, but not to explain it IMO (especially since it's quite complex).

Need to use the testcode directive so that this usage example is tested by doctest (in the CI or manually using uv run make -C ./docs doctest), or no block and prefix each line with >>>

Otherwise the example isn't tested by doctest and we have no guarantee that it works and that it keeps working in future updates.

Same comment for the other usage example, and even for any future example you may write.

Would it be possible to have that inside a skill related to documentation for agents?

Too simple for a skill but we should definitely have that explained in AGENTS.md and / or contributing.md, or even have a script checking that we don't use any code-block directive. Feel free to make a PR with what you think is best.

We need the example to be self-sufficient, so we either need to define loader_1 and loader_2 (which is tedious) or just do something like:

inputs = ...
targets = ...

for i in range(len(input) // 2):
    input_1, input_2 = inputs[2*i], inputs[2*i + 1]
    target_1, target_2 = targets[2*i], targets[2*i + 1]
    ...

With the ... filled with the appriopriate code.

Similar comment for the other example.

To better reflect the new usage, I think we should use G = J1 @ J2.T instead of G = J @ J.T in test_reset_restores_first_step_behavior, test_output_lies_on_simplex, test_update_recurrence and test_changing_m_auto_resets.

Similarly, we should use G1 = J1 @ J2.T and G2 = J3 @ J4.T in test_two_consecutive_steps.

I think we can remove this test.

In the official implementation and in LibMTL, the default value of rho is 0.1. It would be better to match this IMO.

If we follow Equation 2.9b from the paper, this should be

losses = ((losses_1 + losses_2) / 2.0)
losses.backward(weights)

In the official implementation, it's also what they do, except that they forgot the division by 2.

I think there was some confusion on discord when we talked about how to project onto the simplex. We all thought that the official implementation was using a softmax, but it (and LibMTL) actually uses:

    def _projection2simplex(self, y):
        m = len(y)
        sorted_y = torch.sort(y, descending=True)[0]
        tmpsum = 0.0
        tmax_f = (torch.sum(y) - 1.0)/m
        for i in range(m-1):
            tmpsum+= sorted_y[i]
            tmax = (tmpsum - 1)/ (i+1.0)
            if tmax > sorted_y[i+1]:
                tmax_f = tmax
                break
        return torch.max(y - tmax_f, torch.zeros(m).to(y.device))

Should we use this way of projecting @PierreQuinton ?

If we do that, we'll need to say that parts of this file were adapted from the official implementation, add a link to it, and add a notice in NOTICES @KhusPatel4450.

I think I know what happened now, the code that I was told to read was from Rasa's MoCo.py and that used torch.softmax, but yeah now I see that it uses this.

I personally think we should follow this

Maybe we could add that this is the behavior of MoDo in LibMTL and in the official implementation when three_grads is True.

Maybe we could add that this is MoDo as described in the paper, and it's the behavior of the official implementation when three_grads is False.